#include <iostream>
#include <vector>
#include <list>
#include <algorithm>

#include <TROOT.h>
#include <TFile.h>
#include <TTree.h>

#include "GeomSvc.h"
#include "TriggerRoad.h"


int main(int argc, char* argv[])
{
  TriggerRoad* p_road = new TriggerRoad();
  TriggerRoad* m_road = new TriggerRoad();
  
  TFile* dataFile = new TFile(argv[1], "READ");
  TTree* dataTree = (TTree*)dataFile->Get("save");
  
  Double_t mass, xF, x1, x2;
  Double_t weight;
  Double_t x0, y0, z0;

  Double_t px1, py1, pz1, p1;
  Double_t px2, py2, pz2, p2;


  dataTree->SetBranchAddress("p_road", &p_road);
  dataTree->SetBranchAddress("m_road", &m_road);

  dataTree->SetBranchAddress("weight", &weight);
  dataTree->SetBranchAddress("z0", &z0);
  dataTree->SetBranchAddress("mass", &mass);
  dataTree->SetBranchAddress("xF", &xF);

  dataTree->SetBranchAddress("px1", &px1);
  dataTree->SetBranchAddress("py1", &py1);
  dataTree->SetBranchAddress("pz1", &pz1);
  dataTree->SetBranchAddress("px2", &px2);
  dataTree->SetBranchAddress("py2", &py2);
  dataTree->SetBranchAddress("pz2", &pz2);

  ///Read positive and negative roads
  std::list<TriggerRoad> p_roads, m_roads;
  for(int i = 0; i < dataTree->GetEntries(); i++)
    {
      dataTree->GetEntry(i);
            if (mass<4) continue;

      p_roads.push_back(*p_road);
      m_roads.push_back(*m_road);
    
      p_road->clear();
      m_road->clear();
    }

  std::list<TriggerRoad> p_roads_unique; p_roads_unique.clear();
  for(std::list<TriggerRoad>::iterator iter = p_roads.begin(); iter != p_roads.end(); ++iter)
    {
      std::list<TriggerRoad>::iterator road = std::find(p_roads_unique.begin(), p_roads_unique.end(), *iter);
      if(road == p_roads_unique.end())
	{
	  p_roads_unique.push_back(*iter);
	}
      else
	{
	  road->targetWeight += iter->targetWeight;
	  road->dumpWeight += iter->dumpWeight;
	}
    }
  p_roads_unique.sort();

  std::list<TriggerRoad> m_roads_unique; m_roads_unique.clear();
  for(std::list<TriggerRoad>::iterator iter = m_roads.begin(); iter != m_roads.end(); ++iter)
    {
      std::list<TriggerRoad>::iterator road = std::find(m_roads_unique.begin(), m_roads_unique.end(), *iter);
      if(road == m_roads_unique.end())
	{
	  m_roads_unique.push_back(*iter);
	}
      else
	{
	  road->targetWeight += iter->targetWeight;
	  road->dumpWeight += iter->dumpWeight;
	}
    }
  m_roads_unique.sort();

  ///Assign unique road ID
  int uniqueID = 0;
  for(std::list<TriggerRoad>::iterator iter = p_roads_unique.begin(); iter != p_roads_unique.end(); ++iter) 
    {
      iter->roadID = uniqueID;
      ++uniqueID;
    }

  uniqueID = 0;
  for(std::list<TriggerRoad>::iterator iter = m_roads_unique.begin(); iter != m_roads_unique.end(); ++iter) 
    {
      iter->roadID = uniqueID;
      ++uniqueID;
    }

  ///Form dimuon trigger
  std::list<Trigger> dimuons; dimuons.clear();
  for(int i = 0; i < dataTree->GetEntries(); i++)
    {
      dataTree->GetEntry(i);
           if(mass < 4) continue;

      Trigger road_pair;
      std::list<TriggerRoad>::iterator p_iter = std::find(p_roads_unique.begin(), p_roads_unique.end(), *p_road);
      std::list<TriggerRoad>::iterator m_iter = std::find(m_roads_unique.begin(), m_roads_unique.end(), *m_road);

      road_pair.p_road = *p_iter;
      road_pair.m_road = *m_iter;

      road_pair.targetWeight = p_road->targetWeight;
      road_pair.dumpWeight = p_road->dumpWeight;

      dimuons.push_back(road_pair);

      //std::cout << i << ": " << p_iter->roadID << "  " << m_iter->roadID << "  " << road_pair.ratio() << "  " << p_road->targetWeight << "  " << p_road->dumpWeight << std::endl;
    }

  std::list<Trigger> dimuons_unique; dimuons_unique.clear();
  for(std::list<Trigger>::iterator iter = dimuons.begin(); iter != dimuons.end(); ++iter)
    {
      std::list<Trigger>::iterator dimuon = std::find(dimuons_unique.begin(), dimuons_unique.end(), *iter);
      if(dimuon == dimuons_unique.end())
	{
	  dimuons_unique.push_back(*iter);
	}
      else
	{
	  dimuon->targetWeight += iter->targetWeight;
	  dimuon->dumpWeight += iter->dumpWeight;
	}
    }
  dimuons_unique.sort();

  ///Assign unique dimuon trigger ID
  uniqueID = 0;
  for(std::list<Trigger>::iterator iter = dimuons_unique.begin(); iter != dimuons_unique.end(); ++iter)
    {
      iter->triggerID = uniqueID;
      ++uniqueID;
      //std::cout << iter->p_road.roadID << "   " << iter->m_road.roadID << "  " << iter->ratio() << "  " << iter->sum() << std::endl;
    }

  Trigger* trigger = new Trigger();

  TFile* saveFile = new TFile(argv[2], "recreate");

  TTree* saveTree1 = new TTree("dimuon", "dimuon");
  saveTree1->Branch("trigger", &trigger, 256000, 99);
  for(std::list<Trigger>::iterator iter = dimuons_unique.begin(); iter != dimuons_unique.end(); ++iter)
    {
      *trigger = *iter;
      saveTree1->Fill();
    }

  TTree* saveTree2 = new TTree("single_p", "single_p");
  saveTree2->Branch("road", &p_road, 256000, 99);
  for(std::list<TriggerRoad>::iterator iter = p_roads_unique.begin(); iter != p_roads_unique.end(); ++iter)
    {
      *p_road = *iter;
      saveTree2->Fill();
    }

  TTree* saveTree3 = new TTree("single_m", "single_m");
  saveTree3->Branch("road", &m_road, 256000, 99);
  for(std::list<TriggerRoad>::iterator iter = m_roads_unique.begin(); iter != m_roads_unique.end(); ++iter)
    {
      *m_road = *iter;
      saveTree3->Fill();
    }

  saveFile->cd();
  saveTree1->Write();
  saveTree2->Write();
  saveTree3->Write();
  saveFile->Close();

  return 1;
}
